from typing import Any, Dict, List, Union
from ..shard.placement_types import (
    DeviceMesh,Replicate,PlacementStrategy,Shard,DTensorSpec,Shard,OperationData, OperationDataType, ShardingStrategy
)
import copy
from .strategy_generator import StrategyGenerator


class WhereStrategyGenerator(StrategyGenerator):
    """
    EmbeddingStrategyGenerator is a generic class to generate strategies for nn.Embedding or F.embedding.
    The operation data is defined as `output = input x other`.
    """
    def __init__(self,
                 operation_data_mapping,
                 device_mesh,
                 x_strategies_vector):
        super().__init__(operation_data_mapping,device_mesh)
        self.x_strategies_vector = x_strategies_vector
    # def validate(self) -> bool:
    #     return super().validate()

    def collate_strategies(self) -> List[ShardingStrategy]:
        strategy_list = []
        # 复制x的策略

        for stategy in self.x_strategies_vector:
            x_specs = copy.deepcopy(stategy.sharding_specs.output_specs)
            x_specs.tensor_meta = self.op_data['x'].data
            condition_spec = DTensorSpec(self.device_mesh,x_specs.placements,self.op_data['condition'].data)  
            output_specs = DTensorSpec(self.device_mesh,x_specs.placements,self.op_data['output'].data)
            # y_spec = DTensorSpec(self.device_mesh,(Replicate(),),self.op_data['y'].data)
            none_specs = DTensorSpec(self.device_mesh,(Replicate(),),None)
            if self.op_data['x'].data.shape:
                input_specs=[condition_spec,x_specs,none_specs]
            else:
                input_specs=[condition_spec,none_specs,x_specs]
            sharding_specs = PlacementStrategy(input_specs=input_specs,output_specs=output_specs)
            strategy_list.append(ShardingStrategy(name=str(sharding_specs),sharding_specs=sharding_specs,compute_cost=0))
        return strategy_list

